package drds.plus.datasource.api;

import drds.plus.common.jdbc.SqlTypeParser;
import drds.plus.datasource.configuration.DatabaseStatus;
import lombok.extern.slf4j.Slf4j;

import java.sql.SQLException;
import java.sql.Statement;

@Slf4j
public class StatementWrapper {
    protected final DataSourceWrapper datasourceWrapper;
    protected final ConnectionWrapper connectionWrapper;

    final String applicationId;
    public Statement statement;
    /**
     * 经过计算后的结果集，允许使用 getResult函数调用. 一个statement只允许有一个结果集
     */
    protected ResultSetWrapper resultSetWrapper;


    public StatementWrapper(Statement statement, ConnectionWrapper connectionWrapper, DataSourceWrapper tdsWrapper, String applicationId) {
        this.statement = statement;
        this.connectionWrapper = connectionWrapper;
        this.datasourceWrapper = tdsWrapper;

        this.applicationId = applicationId;
    }

    public void addBatch(String sql) throws SQLException {
        this.statement.addBatch(sql);
    }

    public void cancel() throws SQLException {
        this.statement.cancel();
    }

    public void clearBatch() throws SQLException {
        this.statement.clearBatch();
    }

    public void close() throws SQLException {
        close(true);
    }

    void close(boolean removeThis) throws SQLException {
        try {
            if (resultSetWrapper != null) {
                resultSetWrapper.close();
            }
        } catch (SQLException e) {
            log.warn("Close resultSetWrapper failed.", e);
        } finally {
            resultSetWrapper = null;
        }

        try {
            if (this.statement != null)
                this.statement.close();
        } finally {
            this.statement = null; // 端口与物理statement的引用，底下可能会有ps
            // cache，导致节点无法被gc
            if (removeThis) {
                // 关闭之后，移除
                connectionWrapper.removeStatementWrapper(this);
            }
        }
    }

    protected void recordReadTimes() throws SQLException {
        DatabaseStatus dataBaseStatus = datasourceWrapper.dataSourceConnectionProperties.dataBaseStatus;
        if (dataBaseStatus != DatabaseStatus.R_STATUS && dataBaseStatus != DatabaseStatus.RW_STATUS) {
            throw new SQLException("db do not allow to execute read ! dataBaseStatus is " + dataBaseStatus);
        }

        if (!datasourceWrapper.readTimeSliceFlowControl.allow()) {
            throw new SQLException(datasourceWrapper.readTimeSliceFlowControl.reportExceed());
        }
    }

    protected void recordWriteTimes() throws SQLException {
        DatabaseStatus dataBaseStatus = datasourceWrapper.dataSourceConnectionProperties.dataBaseStatus;
        if (dataBaseStatus != DatabaseStatus.W_STATUS && dataBaseStatus != DatabaseStatus.RW_STATUS) {
            throw new SQLException("db do not allow to execute write ! dataBaseStatus is " + dataBaseStatus);
        }

        if (!datasourceWrapper.writeTimeSliceFlowControl.allow()) {
            throw new SQLException(datasourceWrapper.writeTimeSliceFlowControl.reportExceed());
        }
    }

    // 增加并发读计数并判断阀值
    protected void increaseConcurrentRead() throws SQLException {
        int maxConcurrentReadRestrict = datasourceWrapper.dataSourceConnectionProperties.maxConcurrentReadRestrict;
        int concurrentReadCount = datasourceWrapper.concurrentReadCount.incrementAndGet();
        if (maxConcurrentReadRestrict != 0) {
            if (concurrentReadCount > maxConcurrentReadRestrict) {
                datasourceWrapper.readTimesReject.incrementAndGet();
                throw new SQLException("maxConcurrentReadRestrict reached , " + maxConcurrentReadRestrict);
            }
        }
    }

    // 增加并发写计数并判断阀值
    protected void increaseConcurrentWrite() throws SQLException {
        int maxConcurrentWriteRestrict = datasourceWrapper.dataSourceConnectionProperties.maxConcurrentWriteRestrict;
        int concurrentWriteCount = datasourceWrapper.concurrentWriteCount.incrementAndGet();
        if (maxConcurrentWriteRestrict != 0) {
            if (concurrentWriteCount > maxConcurrentWriteRestrict) {
                datasourceWrapper.writeTimesReject.incrementAndGet();
                throw new SQLException("maxConcurrentWriteRestrict reached , " + maxConcurrentWriteRestrict);
            }
        }
    }

    // 减少并发读计数
    protected void decreaseConcurrentRead() throws SQLException {
        datasourceWrapper.concurrentReadCount.decrementAndGet();
    }

    // 减少并发写计数
    protected void decreaseConcurrentWrite() throws SQLException {
        datasourceWrapper.concurrentWriteCount.decrementAndGet();
    }

    public boolean execute(String sql) throws SQLException {
        return executeInternal(sql, -1);
    }

    public boolean execute(String sql, int autoGeneratedKeys) throws SQLException {
        return executeInternal(sql, autoGeneratedKeys);
    }

    private boolean executeInternal(String sql, int autoGeneratedKeys) throws SQLException {
        if (SqlTypeParser.isQuerySql(sql)) {
            executeQuery(sql);
            return true;
        } else {
            executeUpdateInternal(sql, autoGeneratedKeys);
            return false;
        }
    }

    private int executeUpdateInternal(String sql, int autoGeneratedKeys) throws SQLException {

        ensureResultSetIsEmpty();
        recordWriteTimes();
        increaseConcurrentWrite();

        try {

            if (autoGeneratedKeys != -1) {
                return this.statement.executeUpdate(sql, autoGeneratedKeys);
            } else {
                return this.statement.executeUpdate(sql);
            }
        } catch (SQLException e) {

            throw e;
        } finally {

            decreaseConcurrentWrite();
        }
    }

    public int[] executeBatch() throws SQLException {
        ensureResultSetIsEmpty();
        recordWriteTimes();
        increaseConcurrentWrite();
        try {
            return this.statement.executeBatch();
        } finally {
            decreaseConcurrentWrite();
        }
    }

    public ResultSetWrapper executeQuery(String sql) throws SQLException {

        ensureResultSetIsEmpty();
        recordReadTimes();
        increaseConcurrentRead();

        try {
            resultSetWrapper = new ResultSetWrapper(this, this.statement.executeQuery(sql));
            return resultSetWrapper;
        } catch (SQLException e) {
            decreaseConcurrentRead();

            throw e;
        } finally {

        }
    }

    public int executeUpdate(String sql) throws SQLException {
        return executeUpdateInternal(sql, -1);
    }

    public int executeUpdate(String sql, int autoGeneratedKeys) throws SQLException {
        return executeUpdateInternal(sql, autoGeneratedKeys);
    }

    public ConnectionWrapper getConnection() throws SQLException {
        return connectionWrapper;
    }

    public int getFetchSize() throws SQLException {
        return this.statement.getFetchSize();
    }

    public void setFetchSize(int rows) throws SQLException {
        this.statement.setFetchSize(rows);
    }

    public ResultSetWrapper getGeneratedKeys() throws SQLException {
        return new ResultSetWrapper(this, this.statement.getGeneratedKeys());
    }

    public int getMaxFieldSize() throws SQLException {
        return this.statement.getMaxFieldSize();
    }

    public void setMaxFieldSize(int max) throws SQLException {
        this.statement.setMaxFieldSize(max);
    }

    public int getMaxRows() throws SQLException {
        return this.statement.getMaxRows();
    }

    public void setMaxRows(int max) throws SQLException {
        this.statement.setMaxRows(max);
    }

    public boolean getMoreResults() throws SQLException {
        return this.statement.getMoreResults();
    }

    public boolean getMoreResults(int current) throws SQLException {
        return this.statement.getMoreResults(current);
    }

    public int getQueryTimeout() throws SQLException {
        return this.statement.getQueryTimeout();
    }

    public void setQueryTimeout(int seconds) throws SQLException {
        this.statement.setQueryTimeout(seconds);
    }

    public ResultSetWrapper getResultSet() throws SQLException {

        return resultSetWrapper;
    }

    public int getUpdateCount() throws SQLException {
        return this.statement.getUpdateCount();
    }

    /**
     * 如果新建了查询，那么上一次查询的结果集应该被显示的关闭掉。这才是符合jdbc规范的
     */
    protected void ensureResultSetIsEmpty() throws SQLException {
        if (resultSetWrapper != null) {
            try {
                resultSetWrapper.close();
            } catch (SQLException e) {
                log.error("exception on close last result set . can do nothing..", e);
            } finally {
                resultSetWrapper = null;
            }
        }
    }

    public boolean isClosed() throws SQLException {
        return this.statement.isClosed();
    }

    public String getApplicationId() {
        return this.applicationId;
    }
}
